package org.example.kafka.rebalance;

import com.sun.scenario.effect.Offset;
import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;

import java.util.Collection;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 再均衡监听器
 */
public class HandlerRebalance implements ConsumerRebalanceListener {
    // 模拟一个保存分区偏移量的数据库表
    public final  static ConcurrentHashMap<TopicPartition,Long> partitionOffsetMap = new ConcurrentHashMap<>();

    private final Map<TopicPartition, OffsetAndMetadata> currOffsets;
    private final KafkaConsumer<String,String> consumer;
    public HandlerRebalance(Map<TopicPartition, OffsetAndMetadata> currOffsets,
                            KafkaConsumer<String, String> consumer) {
        this.currOffsets = currOffsets;
        this.consumer = consumer;
    }

    /**
     * 分区再平衡之前
     * @param partitions
     */
    @Override
    public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
        final String id = Thread.currentThread().getId()+"";
        System.out.println(id+"-onPartitionsRevoked参数值为："+partitions);
        System.out.println(id+"-服务器准备分区再均衡，提交偏移量。当前偏移量为："
                +currOffsets);
        //偏移量写入数据库
        for (TopicPartition topicPartition:partitions){
            partitionOffsetMap.put(topicPartition, currOffsets.get(topicPartition).offset());
        }
        consumer.commitSync(currOffsets);
        //提交业务数和偏移量入库  tr.commit
    }

    /**
     * 分区再均衡完成以后
     * @param partitions
     */
    @Override
    public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
        final String id = Thread.currentThread().getId()+"";
        System.out.println(id+"-再均衡完成，onPartitionsAssigned参数值为："+partitions);
        System.out.println("分区偏移量表中："+partitionOffsetMap);

        for (TopicPartition topicPartition:partitions){
//            System.out.println(id+"-topicPartition"+topicPartition);
            // 模拟从数据库中取出上次的偏移量
            Long offset = partitionOffsetMap.get(topicPartition);
            if(offset==null) {
                continue;
            }
            // kafka提供了seek方法，可以让我们从分区的固定位置开始消费。
            consumer.seek(topicPartition,partitionOffsetMap.get(topicPartition));
        }
    }
}
